#!/usr/bin/env python3
"""Compute Welch t-test for monthly brightness grouped by border_city.

Reads `data/云南乡镇月度灯光亮度.parquet` using duckdb, computes n/mean/sd per group,
then runs scipy.stats.ttest_ind_from_stats and computes Cohen's d.
Prints formatted results with 2 decimal places.
"""
import math
import duckdb
from scipy import stats

PARQUET = 'data/云南乡镇月度灯光亮度.parquet'

q = f"""
SELECT border_city, COUNT(*) AS n, AVG(日度夜间灯光亮度) AS mean, STDDEV_SAMP(日度夜间灯光亮度) AS sd
FROM read_parquet('{PARQUET}')
GROUP BY border_city
ORDER BY border_city;
"""

con = duckdb.connect(database=':memory:')
res = con.execute(q).fetchall()
# Expect two rows border_city 0 and 1
if len(res) < 2:
    raise SystemExit('Expected at least two groups for border_city')

rows = {row[0]: {'n': int(row[1]), 'mean': float(row[2]), 'sd': float(row[3])} for row in res}

# ensure we have 0 and 1 keys
if 0 not in rows or 1 not in rows:
    raise SystemExit('Expected border_city values 0 and 1')

g0 = rows[0]
g1 = rows[1]

# Welch t-test from stats
tstat, pvalue = stats.ttest_ind_from_stats(mean1=g0['mean'], std1=g0['sd'], nobs1=g0['n'],
                                           mean2=g1['mean'], std2=g1['sd'], nobs2=g1['n'],
                                           equal_var=False)

# Cohen's d (using pooled sd approximation for unequal n and sd: sqrt(((sd1^2 + sd2^2)/2)))
pooled_sd = math.sqrt((g0['sd'] ** 2 + g1['sd'] ** 2) / 2)
cohen_d = (g0['mean'] - g1['mean']) / pooled_sd if pooled_sd != 0 else float('inf')

print('group | n | mean | sd')
print('-----------------------')
print(f"0     | {g0['n']:,} | {g0['mean']:.2f} | {g0['sd']:.2f}")
print(f"1     | {g1['n']:,} | {g1['mean']:.2f} | {g1['sd']:.2f}")
print('\nWelch t-test (two-sided):')
print(f"t = {tstat:.2f}, p = {pvalue:.2e}")
print(f"Cohen's d = {cohen_d:.2f}")
